from random import random
from networkx import complete_graph, cycle_graph, wheel_graph, star_graph
import scipy.stats.distributions as distributions
from enthought.traits.api import HasPrivateTraits, Int, Range, Enum, Bool, Any, Disallow

class World(HasPrivateTraits):
    """
    Base class defining the world in which an IDS game takes place, including
    the network type, payoff structure, etc.
    
    This class may be overrided to create alternative worlds. 
    """
    start_balance = Int(0)
    cost_of_investing = Range(low=-100.0, high=0.0, value=-12.0)
    cost_of_negevent = Range(low=-100.0, high=0.0, value=-50.0)
    p = Range(low=0.0, high=1.0, value=0.2, label='Prob of negevent')
    graph_type = Enum(['complete', 'cycle', 'star', 'wheel'], label='Graph')
    num_rounds = Range(low=1, high=100, value=10, label='Rounds/game')
    correlated_outcomes = Bool(False)
    divide_cost = Bool(False)

    def __init__(self,  start_balance=0, cost_of_investing=-12, 
            cost_of_negevent=-50, p=.2, graph_type='complete',
            num_rounds=10, correlated_outcomes=False, 
            divide_cost=False, **traits):
        """
        Create a new IDS world.
        
        Keyword arguments:
        start_balance -- starting balance for players.
        cost_of_investing -- cost to invest in each round (must be neg)
        cost_of_negevent -- loss when there is a negative event (must be neg)
        p -- independent prob. of a negative event originating from any agent
        graph_type -- network structure (e.g., 'complete', 'cycle')
        num_rounds -- the number of rounds in a game
        correlated_outcomes -- compute outcomes separately for each agent, 
            rather than correlating outcomes among neighbors
        divide_cost -- divide cost of a negative outcome among neighbors, 
            rather than having each agent pay its full amount. (not yet 
            implemented)
        """
        # initaliaze superclass
        HasPrivateTraits.__init__(self, **traits)
        self._graph = None

    def setup_agents(self, agents):
        """Takes a set of agent instances and places them in a new graph."""
        self._graph = self._init_graph(len(agents),self.graph_type)
        
    def _init_graph(self, num_agents, graph_type):
        """Generate network graph"""
        # uses networkx for graph creation (consult networkx documentation for
        #  other default graph types (e.g., random graphs)
        if graph_type == 'complete':
            # complete graph, i.e., every agent is connected to every other 
            return complete_graph(num_agents)
        if graph_type == 'cycle':
            # cycle graph, i.e., agents are in a ring
            return cycle_graph(num_agents)
        if graph_type == 'wheel':
            # wheel graph, i.e., agents are in a ring with a hub
            return wheel_graph(num_agents)
        if graph_type == 'star':
            # star graph, i.e., one agent connected to all others
            return star_graph(num_agents-1)
        else:
            raise ValueError("Graph type %s not recognized"%graph_type)

    def next_round(self):
        """
        Called by game class after each round is played.
        May be used, for example, to alter the graph between rounds.
        """
        assert(self._graph != None)
        # not implemented

    def neighbors(self,agent_i):
        """Returns a list of neighbors for the given agent number."""
        assert(self._graph != None)
        return self._graph[agent_i].keys()

    def num_invested(self,agent_i,decisions):
        """
        Given an agent id and a true or hypothetical list of invest decisions, 
        return the number of agents in the given agent's neighborhood that 
        have invested (inclusive of the given agent).
        
        Keyword arguments:
        agent_i -- id of an agent
        decisions -- list of invest decisions by all agents, ordered by agent id
        """
        assert(self._graph != None and len(decisions) == len(self._graph))
        return sum([decisions[i] for i in (self.neighbors(agent_i)+[agent_i])])

    def p_negevent(self, agent_i, decisions):
        """
        Given an agent id and a true or hypothetical list of invest decisions, 
        return the probability that a negative event affects the given agent
        
        Keyword arguments:
        agent_i -- id of an agent
        decisions -- list of invest decisions by all agents, ordered by agent id
        """
        assert(self._graph != None and len(decisions) == len(self._graph))
        neighborhood_size = len(self.neighbors(agent_i))+1
        num_invested = self.num_invested(agent_i, decisions)
        # We use a binomial distribution and compute expectation of at least
        # 1 event occuring (i.e., 1 minus the probability of 0 events occuring)
        return 1 - distributions.binom.pmf(0,neighborhood_size-num_invested,self.p)         

    def outcomes(self, decisions):
        """
        Given list of decisions of all agents, return list of outcomes for 
        this round.
        """
        assert(self._graph != None and len(decisions) == len(self._graph))
        ids = range(len(decisions))
        if self.correlated_outcomes:
            # generate an independent outcome for each agent 
            o = [(random() < (1-x) * self.p) for x in decisions]
             # spread negative outcomes to neighbors
            negevent = [any([o[n] for n in self.neighbors(i)+[i]]) for i in ids]
        else:
            # compute probability of negative events for each agent
            p_negevent = [self.p_negevent(i,decisions) for i in ids]
            # generate an independent outcome for each agent
            negevent = [random() < x for x in p_negevent]
        # compute outcomes/payoffs for each agent
        outcome = [(decisions[i] * self.cost_of_investing) + \
                   (negevent[i] * self.cost_of_negevent) for i in ids]
        return outcome,negevent
    
    def parameters_as_str(self):
        """Return a string with all parameters."""
        keys = filter(lambda x:x[0] != '_', sorted(self.trait_get().keys()))
        result = ','.join(keys) + '\n'
        result += ','.join([str(self.trait_get()[k]) for k in keys]) + '\n'
        return result